from pathlib import Path
import pathlib
from ruamel.yaml import YAML
from copy import deepcopy
from treelib import Tree
import os
import argparse
import re

"""
auto 文件夹下存放着一些脚本，用来统一执行
python auto_gen_dockercp.py --aarch64
"""
parser = argparse.ArgumentParser()
parser.add_argument("--aarch64", action="store_true", default="", help="生成 aarch64 的 docker-compose 以及对应的 Dockerfiles")
args = parser.parse_args()

image_base_name = "harbor.x-contion.top:500/x-contion"
public_list_image_base_name = ["swr.cn-north-4.myhuaweicloud.com/wlf_darkmaster/x-contion"]


yaml = YAML(typ="safe")

print(__file__)
path_root = Path(__file__).parent

dict_services = {}
yaml_dockercp = {"services": dict_services}
yaml_server = {
    "environment": [
        "SHOW=true",
        # "DOCKER_BUILDKIT=1",
        "BUILDKIT_TTY_LOG_LINES=18",
    ],
    "build": {
        "network": "host",
        "context": ".",
        "dockerfile": "none",
        "args": [
            "BUILDKIT_PROGRESS=plain",
        ],
    },  #! 注意替换
    "image": "none",  #! 注意替换
    "hostname": "ubuntu",
    # "entrypoint": "/usr/sbin/sshd -D &",
}

#! 构建时准从这个列表的优先级进行构建

tree_core = Tree()
tree_core.create_node("base_ubuntu", "base_ubuntu")
tree_core.create_node("base_compile", "base_compile", parent="base_ubuntu")
tree_core.create_node("base_python", "base_python", parent="base_ubuntu")

tree_core.create_node("steam", "steam", parent="base_ubuntu")
tree_core.create_node("lib_opencv", "lib_opencv", parent="base_compile")

# tree_core.create_node("ascend_cann", "ascend_cann", parent="base_python") #! CANN 直接在要用的时候去用就行了，不用单独弄一个镜像
tree_core.create_node("pytorch_cuda", "pytorch_cuda", parent="base_python")
tree_core.create_node("pytorch_npu", "pytorch_npu", parent="base_python")
tree_core.create_node("rl_npu", "rl_npu", parent="pytorch_npu")

list_tree = [tree_core]

def generate_compose_yaml(platform="amd64"):
    path_dockerfiles_cp = None
    if platform == "aarch64":
        print("生成 aarch64 的 docker-compose 以及对应的 Dockerfiles")
        path_dockerfiles_cp = path_root.joinpath("dockerfiles_aarch64")
        path_dockerfiles_cp.mkdir(exist_ok=True)
    pass

    for the_dir in path_root.joinpath("dockerfiles").glob("*"):
        print(the_dir.name)
        if path_dockerfiles_cp is not None:
            path_cp_platform = path_dockerfiles_cp.joinpath(the_dir.name)
            path_cp_platform.mkdir(exist_ok=True)

        for path_dockerfile in the_dir.glob("Dockerfile.*"):

            if path_dockerfiles_cp is not None:
                path_cp = path_cp_platform.joinpath(path_dockerfile.name)
                with open(path_dockerfile, "r") as f:
                    content = f.readlines()
                for _i, line_i in enumerate(content):
                    if line_i.startswith("FROM "):
                        content[_i]= line_i.replace(image_base_name, image_base_name + "-aarch64")
                        continue #* 多阶段构建需要替换多个 FROM

                with open(path_cp, "w") as f:
                    f.writelines(content)
                path_dockerfile = path_cp
            server_name, server_info = input_info(path_dockerfile, platform)
            dict_services[server_name] = server_info
            # print(dockerfile)
    yaml.dump(yaml_dockercp, path_root.joinpath("./docker-compose.yaml"))


def generate_build_script(dict_services, path_script):
    # * 需要处理依赖关系
    dict_imagename = {}
    dict_image_depand = {}
    for _, value in dict_services.items():
        path_dockerfile = value["build"]["dockerfile"]
        image_name = value["image"]
        dict_imagename[image_name] = path_dockerfile
        dict_image_depand[image_name] = []

    # * 二次循环，判断依赖
    list_image_other = []
    for image_name, path_dockerfile in dict_imagename.items():
        with open(path_root / path_dockerfile, "r") as f:
            content = f.readlines()
            the_list_depand = []
            #
            for line in content:
                line = line.strip()
                if line.startswith("FROM"):
                    depand_image_name = line.split(" ")[1]
                    if depand_image_name in dict_imagename:
                        the_list_depand.append(depand_image_name)
                    else:
                        if depand_image_name not in list_image_other:
                            list_image_other.append(depand_image_name)
            #
            dict_image_depand[image_name] = the_list_depand

    # * 这里是一个比较傻的方法实现依赖顺序排列
    list_service = list(dict_imagename.keys())
    list_build = []
    while len(list_service) > 0:
        target = list_service.pop(0)
        the_list_depand = dict_image_depand[target]
        if len(the_list_depand) == 0:
            list_build.append(target)
            continue
        else:
            for depand in the_list_depand:
                if depand not in list_build:
                    list_service.append(target)
                    break
            else:
                list_build.append(target)

    return list_build, list_image_other


def input_info(path_dockerfile: Path, platform="amd64"):
    """快速写入信息函数

    Args:
        dockerfile (str): dockerfile的路径，以所在文件夹的名称以及文件名标定server_name
        image (str): _description_
        tag (str): _description_
        hostname (str, optional): _description_. Defaults to "ubuntu".
    """
    tag = "default"
    image = path_dockerfile.parent.name
    tags = path_dockerfile.name.split(".")  #! 服务名称以Dockerfile.小数点后面的名称命名,
    tag = tags[1]
    server_name = image + "." + tag
    if len(tags) > 2:
        raise NameError("Dockerfile的名称只能有一个小数点")

    server_info = deepcopy(yaml_server)
    if platform == "aarch64":
        image = image_base_name + "-aarch64/" + image + ":" + tag
    else:
        image = image_base_name + "/" + image + ":" + tag
    server_info["image"] = image
    server_info["build"]["dockerfile"] = path_dockerfile.relative_to(path_root).__str__()

    return server_name, server_info


def generate_build_all(list_build, list_image_other, path_shell):
    dict_services = yaml_dockercp["services"]
    with open(path_shell, "w") as f_shell:
        f_shell.write(r"#!/bin/sh")
        f_shell.write("\n")
        f_shell.write(r'shell_pwd=$(dirname "$(readlink -f "$0")")')
        f_shell.write("\n")
        f_shell.write(r'echo 当前目录为:"$shell_pwd"')
        f_shell.write("\n")
        f_shell.write(r'cd "${shell_pwd}"')
        f_shell.write("\n")
        #! 这一行bash代码用于获取当前local镜像名称,暂时没用上，本来想弄一个只针对未构建过的镜像的功能，但是有点麻烦
        f_shell.write("\n#=================================\n#!=======  REQUIRED IMAGE  =========\n#=================================\n\n")
        for pull_image in list_image_other:
            f_shell.write(f"docker pull {pull_image} &\n")
        f_shell.write("wait\n\n")

        for target in list_build:
            service_name = target.split("/")[-1].replace(":", ".")

            f_shell.write(f"# build {target}\n")
            f_shell.write(f"# {Path(__file__).parent.joinpath(dict_services[service_name]['build']['dockerfile']).__str__()}\n")
            f_shell.write(f"echo -e '\\033[32mdocker-compose build {service_name}\\033[0m'\n")
            f_shell.write("docker-compose build " + service_name + "\n")
            for public_url in public_list_image_base_name:
                f_shell.write(f"# skopeo copy docker-daemon:{target} docker://{target.replace(image_base_name, public_url)}/ &\n")

            f_shell.write(f"docker push {target} &\n")
            pass
        f_shell.write("\n\nwait\n\n")
    os.system(f"chmod +x {path_shell}")


if __name__ == "__main__":
    list_dir = list(path_root.joinpath("dockerfiles").glob("*"))
    args = parser.parse_args()

    if args.aarch64:
        generate_compose_yaml(platform="aarch64")
    else:
        generate_compose_yaml()

    list_build, list_image_other = generate_build_script(dict_services, path_root.joinpath("build_all.sh"))

    generate_build_all(list_build, list_image_other, path_root.joinpath("build_all.sh"))

    print("\033[32m生成完毕\033[0m")
